public class FindDiagonalOrder {

    public int[] findDiagonalOrder(int[][] mat) {
        int[] ans = new int[mat.length * mat[0].length];
        int temp = 0;
        boolean flag = true;
        for (int i = 0; i < mat.length; i++) {
            temp = process(mat, temp, ans, flag, i, 0);
            flag = !flag;
        }
        for (int i = 1; i < mat[0].length; i++) {
            temp = process(mat, temp, ans, flag, mat.length - 1, i);
            flag = !flag;
        }
        return ans;
    }

    public int process(int[][] mat, int i, int[] ans, boolean flag, int m, int n) {
        int index = i;
        while (m < mat.length && m >= 0 && n < mat[0].length && n >=0) {
            ans[i++] = mat[m][n];
            m--;
            n++;
        }
        int res = i;
        if (!flag) {
            i -= 1;
            while (index <= i) {
                int temp = ans[index];
                ans[index] = ans[i];
                ans[i] = temp;
                ++index;
                --i;
            }
        }
        return res;
    }
}
